#!/bin/bash

# -------------------------------------------------------------------------------
# Filename:     run_web_demo.sh
# Update:       2025/04/16
# Description:  启动 Qwen2-5-Omni-7B 的 Web Demo
# Example:      ./run_web_demo.sh [--port 7860] [--ckpt /path/to/checkpoint] [--devices 0,1]
# Depends:      Python, Torch2.5 环境, web_demo.py
# -------------------------------------------------------------------------------

set -e

# 默认配置
DEFAULT_PORT=7860
DEFAULT_CKPT_PATH="/data/models/llm/models/Qwen2.5-Omni-7B"
DEFAULT_SERVER_IP="0.0.0.0"
DEFAULT_UI_LANG="zh"
DEFAULT_DEVICES="0,1,2,3"
SCRIPT_DIR="/home/share/pytorch2.5/qwen2.5-omni/tools"
PY_SCRIPT="web_demo.py"
LOGFILE="web_demo.log"

# 初始化参数为默认值
PORT="$DEFAULT_PORT"
CKPT_PATH="$DEFAULT_CKPT_PATH"
SERVER_IP="$DEFAULT_SERVER_IP"
UI_LANG="$DEFAULT_UI_LANG"
DEVICES="$DEFAULT_DEVICES"

# 打印帮助函数
usage() {
  echo ""
  echo "用法: ./run_web_demo.sh [选项]"
  echo ""
  echo "选项:"
  echo "  --port <端口号>           设置 Web 服务端口 (默认: $DEFAULT_PORT)"
  echo "  --ckpt <路径>             模型 checkpoint 路径 (默认: $DEFAULT_CKPT_PATH)"
  echo "  --lang <zh|en>            UI 显示语言 (默认: $DEFAULT_UI_LANG)"
  echo "  --devices <设备列表>      设置 MLU_VISIBLE_DEVICES (默认: $DEFAULT_DEVICES)"
  echo "  -h, --help                显示此帮助信息并退出"
  echo ""
  echo "示例:"
  echo "  ./run_web_demo.sh --port 8888 --ckpt /data/models/Qwen2.5-Omni-7B --devices 0,2"
  echo ""
  exit 0
}

# 解析命令行参数
while [[ "$#" -gt 0 ]]; do
  case $1 in
    --port) PORT="$2"; shift ;;
    --ckpt) CKPT_PATH="$2"; shift ;;
    --lang) UI_LANG="$2"; shift ;;
    --devices) DEVICES="$2"; shift ;;
    -h|--help) usage ;;
    *) echo "❌ 未知参数: $1"; usage ;;
  esac
  shift
done

# 展示实际使用的参数（含默认标识）
echo "-----------------------------"
echo "🚀 启动 Qwen2.5 Web Demo"
echo "📦 模型路径: $CKPT_PATH" $( [ "$CKPT_PATH" = "$DEFAULT_CKPT_PATH" ] && echo "(默认)" )
echo "🌐 监听地址: $SERVER_IP"
echo "🔌 监听端口: $PORT" $( [ "$PORT" = "$DEFAULT_PORT" ] && echo "(默认)" )
echo "🗣️ UI语言: $UI_LANG" $( [ "$UI_LANG" = "$DEFAULT_UI_LANG" ] && echo "(默认)" )
echo "💻 可见设备: $DEVICES" $( [ "$DEVICES" = "$DEFAULT_DEVICES" ] && echo "(默认)" )
echo "📄 日志文件: $LOGFILE"
echo "🔗 访问链接: http://$SERVER_IP:$PORT"
echo "-----------------------------"

# 检查依赖
command -v python >/dev/null 2>&1 || { echo >&2 "❌ Python 未安装，请先安装 Python。"; exit 1; }
if [ ! -f "$SCRIPT_DIR/$PY_SCRIPT" ]; then
  echo "❌ 找不到 web_demo.py 脚本: $SCRIPT_DIR/$PY_SCRIPT"
  exit 1
fi

# 激活环境
echo "✅ 激活 PyTorch 推理环境..."
source /torch/venv3/pytorch_infer/bin/activate

# 切换目录
cd "$SCRIPT_DIR"

# 设置设备
export MLU_VISIBLE_DEVICES="$DEVICES"

# 启动服务
echo "🚀 正在启动 Web Demo..."
python "$PY_SCRIPT" \
  --flash-attn2 \
  --server-name "$SERVER_IP" \
  --server-port "$PORT" \
  --ui-language "$UI_LANG" \
  --checkpoint-path "$CKPT_PATH" \
  | tee "$LOGFILE"
